import math
import torch
import numpy as np
import torch.nn as nn


d_model = 512  # Embedding Size
d_ff = 2048  # FeedForward dimension   前馈神经网络
d_k = d_v = 64  # dimension of K(=Q), V
n_layers = 6  # number of Encoder of Decoder Layer  # 编码器层数
n_heads = 8  # number of heads in Multi-Head Attention  多头注意力的头数   # 一般dk*nhead要等于dmodel

# 编码，位置嵌入
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.3, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)      # 随机失活

        pe = torch.zeros(max_len, d_model)    #返回一个由标量值0填充的张量，其形状由变量参数size定义
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)   # unsqueeze(1)在第二维增加一个维度
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))   # y = e**x
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)   # unsqueeze(0)在第一维前增加一个维度； torch.transpose(Tensor,dim0,dim1)对ndarray矩阵进行转置的操作
        self.register_buffer('pe', pe)   # 向模块添加持久缓冲区。将tensor注册成buffer,在forward中使用self.mybuffer,而不是self.mybuffer_tmp

    def forward(self, x):
        '''
        x: [seq_len, batch_size, d_model]
        '''
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)


# 注意力机制1
class ScaledDotProductAttention(nn.Module):
    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()

    def forward(self, Q, K, V):
        '''
        Q: [batch_size, n_heads, len_q, d_k]
        K: [batch_size, n_heads, len_k, d_k]
        V: [batch_size, n_heads, len_v(=len_k), d_v]
        attn_mask: [batch_size, n_heads, seq_len, seq_len]
        '''
        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k)  # scores : [batch_size, n_heads, len_q, len_k]
        # torch.transpose(Tensor,dim0,dim1)对ndarray矩阵进行转置的操作
        # torch.matmul()。tensor的乘法：
        # （1）当输入是都是二维时，就是普通的矩阵乘法，和tensor.mm函数用法相同。
        # （2）当输入有多维时，把多出的一维作为batch提出来，其他部分做矩阵乘法。

        attn = nn.Softmax(dim=-1)(scores)  # 归一化指数函数
        context = torch.matmul(attn, V)  # [batch_size, n_heads, len_q, d_v]  # torch.matmul()。tensor的乘法
        return context, attn

# 注意力机制2
class MultiHeadAttention(nn.Module):
    def __init__(self):
        super(MultiHeadAttention, self).__init__()
        self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=False)    # 全连接层   # 一般dk*nhead要等于dmodel
        self.W_K = nn.Linear(d_model, d_k * n_heads, bias=False)
        self.W_V = nn.Linear(d_model, d_v * n_heads, bias=False)
        self.fc = nn.Linear(n_heads * d_v, d_model, bias=False)

    def forward(self, input_Q, input_K, input_V):
        '''
        input_Q: [batch_size, len_q, d_model]
        input_K: [batch_size, len_k, d_model]
        input_V: [batch_size, len_v(=len_k), d_model]
        attn_mask: [batch_size, seq_len, seq_len]
        '''
        residual, batch_size = input_Q, input_Q.size(0)
        # (B, S, D) -proj-> (B, S, D_new) -split-> (B, S, H, W) -trans-> (B, H, S, W)
        Q = self.W_Q(input_Q).view(batch_size, -1, n_heads, d_k).transpose(1, 2)  # Q: [batch_size, n_heads, len_q, d_k]
        K = self.W_K(input_K).view(batch_size, -1, n_heads, d_k).transpose(1, 2)  # K: [batch_size, n_heads, len_k, d_k]
        V = self.W_V(input_V).view(batch_size, -1, n_heads, d_v).transpose(1, 2)  # V: [batch_size, n_heads, len_v(=len_k), d_v]
        # 在torch里面，view函数相当于numpy的reshape

        # context: [batch_size, n_heads, len_q, d_v], attn: [batch_size, n_heads, len_q, len_k]
        context, attn = ScaledDotProductAttention()(Q, K, V)    # 注意力机制1
        context = context.transpose(1, 2).reshape(batch_size, -1, n_heads * d_v)
        # context: [batch_size, len_q, n_heads * d_v]
        output = self.fc(context)  # [batch_size, len_q, d_model]
        return nn.LayerNorm(d_model).cuda()(output + residual), attn

# 前馈神经网络
class PoswiseFeedForwardNet(nn.Module):
    def __init__(self):
        super(PoswiseFeedForwardNet, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(d_model, d_ff, bias=False),      #全连接层
            nn.ReLU(),    #激活函数
            nn.Linear(d_ff, d_model, bias=False)
        )

    def forward(self, inputs):
        '''
        inputs: [batch_size, seq_len, d_model]
        '''
        residual = inputs
        output = self.fc(inputs)    # def __init(self):里的fc
        return nn.LayerNorm(d_model).cuda()(output + residual)  # [batch_size, seq_len, d_model]
        # nn.LayerNorm()归一化层 对比BN

# 编码器层
class EncoderLayer(nn.Module):
    def __init__(self):
        super(EncoderLayer, self).__init__()
        self.enc_self_attn = MultiHeadAttention()   # 注意力机制2
        self.pos_ffn = PoswiseFeedForwardNet()     # 前馈神经网络

    def forward(self, enc_inputs):
        '''
        enc_inputs: [batch_size, src_len, d_model]
        enc_self_attn_mask: [batch_size, src_len, src_len]
        '''
        enc_outputs, attn = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs, )  # enc_inputs to same Q,K,V  # 注意力机制2
        enc_outputs = self.pos_ffn(enc_outputs)  # enc_outputs: [batch_size, src_len, d_model]   # 前馈神经网络
        return enc_outputs, attn

# 编码器
class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.pos_emb = PositionalEncoding(d_model)       # # 编码，位置嵌入
        self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)])
        # nn.ModuleList，它是一个储存不同 module，并自动将每个 module 的 parameters 添加到网络之中的容器。
        # # 编码器层

    def forward(self, enc_inputs):
        '''
        enc_inputs: [batch_size, src_len]
        '''
        enc_outputs = enc_inputs
        enc_outputs = self.pos_emb(enc_outputs.transpose(0, 1)).transpose(0, 1)  # [batch_size, src_len, d_model]
        # 转置
        enc_self_attns = []
        for layer in self.layers:
            enc_outputs, enc_self_attn = layer(enc_outputs)
            enc_self_attns.append(enc_self_attn)
        return enc_outputs, enc_self_attns

# 全连接
class Twofc(nn.Module):
    def __init__(self):
        super(Twofc, self).__init__()
        self.linear = nn.Sequential(
            nn.Linear(3584, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Linear(64, 1),
        )

    def forward(self, x):
        out = self.linear(x)
        return out


class Transformer(nn.Module):
    def __init__(self):
        super(Transformer, self).__init__()
        self.encoder = Encoder().cuda()   # 编码器
        self.projection = Twofc()      # 全连接

    def forward(self, enc_inputs):
        '''
        enc_inputs: [batch_size, src_len]
        dec_inputs: [batch_size, tgt_len]
        '''
        enc_outputs, enc_self_attns = self.encoder(enc_inputs)   # 编码器
        enc_outputs = enc_outputs.reshape(enc_outputs.size(0), -1)
        enc_outputs = self.projection(enc_outputs)

        return enc_outputs
